import os

import cv2
import numpy


def image_read(path, image_size):
    label = 0
    names = []
    training_images, training_labels = [], []

    for dirpath, dirnames, filenames in os.walk(path):
        for dirname in dirnames:
            subject_name = os.path.join(dirpath, dirname)
            names.append(dirname)
            for name in os.listdir(subject_name):
                file = os.path.join(subject_name, name)
                img = cv2.imread(file, cv2.IMREAD_GRAYSCALE)
                if img is None:
                    continue
                img = cv2.resize(img, image_size)
                training_images.append(img)
                training_labels.append(label)
            label += 1
    training_images = numpy.asarray(training_images, numpy.uint8)
    training_labels = numpy.asarray(training_labels, numpy.int32)
    return names, training_images, training_labels


def get_image():
    face_cascades = cv2.CascadeClassifier('./cascade/haarcascade_frontalface_default.xml')

    camera = cv2.VideoCapture(0)
    while cv2.waitKey(1) == -1:
        success, frame = camera.read()
        if success:
            gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
            faces = face_cascades.detectMultiScale(gray, 1.3, 5)
            for (x, y, w, h) in faces:
                cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
                gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
                roi_gray = gray[y:y + h, x:x + w]
                if roi_gray.size == 0:
                    continue
                roi_gray = cv2.resize(roi_gray, training_image_size)

                label, confidence = model.predict(roi_gray)
                text = '%s, confidence =%.2f' % (names[label], confidence)
                cv2.putText(frame, text, (x, y - 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
            cv2.imshow('FAce Recognition', frame)


if __name__ == '__main__':
    path_to_training_images = './data'
    training_image_size = (200, 200)
    names, training_images, training_label = image_read(path_to_training_images, training_image_size)

    model = cv2.face.EigenFaceRecognizer_create()
    model.train(training_images, training_label)
    face_cascades = cv2.CascadeClassifier('./cascades/haarcascade_frontalface_default.xml')


    camera = cv2.VideoCapture(0)
    while cv2.waitKey(1) == -1:
        success, frame = camera.read()
        if success:
            gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
            faces = face_cascades.detectMultiScale(gray, 1.03, 5)
            for (x, y, w, h) in faces:
                cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
                gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
                roi_gray = gray[y:y + h, x:x + w]
                if roi_gray.size == 0:
                    continue
                roi_gray = cv2.resize(roi_gray, training_image_size)

                label, confidence = model.predict(roi_gray)
                text = '%s, confidence =%.2f' % (names[label], confidence)
                cv2.putText(frame, text, (x, y - 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
            cv2.imshow('FAce Recognition', frame)
